import os

from dotenv import load_dotenv
from langchain.agents import create_tool_calling_agent, tool, AgentExecutor
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
from pydantic import BaseModel, Field

from tools.langchain_weather import get_weather

class WriteQuery(BaseModel):
    content: str = Field(description="需要写入文档的具体内容")
@tool(args_schema=WriteQuery)
def write_file(content):
    """
    将指定的内容写入本地文件
    :param content: 必要参数，字符串类型,用于表示需要写入文档的具体内容
    :return：是否成功写入
    """
    if not content:
        return "内容不能为空"
    file_path = "output.txt"
    with open(file_path, "w") as f:
        f.write(content)
    return "已经成功写入本地文件"


if __name__ == '__main__':
    load_dotenv(override=True)
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = ChatOllama(model="qwen3:30b", base_url="http://192.168.97.217:11434")
    tools = [get_weather, write_file]
    prompt = ChatPromptTemplate.from_messages([
        ("system",
         "你是天气助手，请根据用户输入给出相应的天气信息，如果用户需要将查询结果把写入文件"),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])
    # 直接创建代理
    agent = create_tool_calling_agent(model, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    response = agent_executor.invoke({"input": "请告诉我北京和南京天气哪个更热,并保存查询结果"})
    print(response)
